#include "OSCAes.h"
#include "OSCBase64.h"

#define DEF_AES_ECODE_KEY "LeBxDt5Rs8gssAw8"
#define DEF_AES_ECODE_IV "hHsca9N2Jie8BmWR"

/*
 * 32-bit integer manipulation macros (little endian)
 */
#ifndef GET_ULONG_LE
#define GET_ULONG_LE(n, b, i) \
    { \
        (n) = ((unsigned long)(b)[(i)]) | ((unsigned long)(b)[(i) + 1] << 8) | ((unsigned long)(b)[(i) + 2] << 16) | ((unsigned long)(b)[(i) + 3] << 24); \
    }
#endif

#ifndef PUT_ULONG_LE
#define PUT_ULONG_LE(n, b, i)                      \
    {                                              \
        (b)[(i)] = (unsigned char)((n));           \
        (b)[(i) + 1] = (unsigned char)((n) >> 8);  \
        (b)[(i) + 2] = (unsigned char)((n) >> 16); \
        (b)[(i) + 3] = (unsigned char)((n) >> 24); \
    }
#endif

/*
 * Forward S-box & tables
 */
static unsigned char FSb[256];
static unsigned long FT0[256];
static unsigned long FT1[256];
static unsigned long FT2[256];
static unsigned long FT3[256];

/*
 * Reverse S-box & tables
 */
static unsigned char RSb[256];
static unsigned long RT0[256];
static unsigned long RT1[256];
static unsigned long RT2[256];
static unsigned long RT3[256];

/*
 * Round constants
 */
static unsigned long RCON[10];

/*
 * Tables generation code
 */
#define ROTL8(x) ((x << 8) & 0xFFFFFFFF) | (x >> 24)
#define XTIME(x) ((x << 1) ^ ((x & 0x80) ? 0x1B : 0x00))
#define MUL(x, y) ((x && y) ? pow[(log[x] + log[y]) % 255] : 0)

static int aes_init_done = 0;

static void OscAes_Gen_Tables(void)
{
    int i, x, y, z;
    int *pow, *log;
    int *tmp;

    tmp = (int *)malloc((POW_LEN + LOG_LEN) * sizeof(pow[0]));
    if (NULL == tmp)
    {
        printf("Aes Gen Tables failed\n");
        return;
    }

    pow = tmp;
    log = &tmp[POW_LEN];

    /*
     * compute pow and log tables over GF(2^8)
     */
    for (i = 0, x = 1; i < 256; i++)
    {
        pow[i] = x;
        log[x] = i;
        x = (x ^ XTIME(x)) & 0xFF;
    }

    /*
     * calculate the round constants
     */
    for (i = 0, x = 1; i < 10; i++)
    {
        RCON[i] = (unsigned long)x;
        x = XTIME(x) & 0xFF;
    }

    /*
     * generate the forward and reverse S-boxes
     */
    FSb[0x00] = 0x63;
    RSb[0x63] = 0x00;

    for (i = 1; i < 256; i++)
    {
        x = pow[255 - log[i]];

        y = x;
        y = ((y << 1) | (y >> 7)) & 0xFF;
        x ^= y;
        y = ((y << 1) | (y >> 7)) & 0xFF;
        x ^= y;
        y = ((y << 1) | (y >> 7)) & 0xFF;
        x ^= y;
        y = ((y << 1) | (y >> 7)) & 0xFF;
        x ^= y ^ 0x63;

        FSb[i] = (unsigned char)x;
        RSb[x] = (unsigned char)i;
    }

    /*
     * generate the forward and reverse tables
     */
    for (i = 0; i < 256; i++)
    {
        x = FSb[i];
        y = XTIME(x) & 0xFF;
        z = (y ^ x) & 0xFF;

        FT0[i] = ((unsigned long)y) ^
                 ((unsigned long)x << 8) ^
                 ((unsigned long)x << 16) ^ ((unsigned long)z << 24);

        FT1[i] = ROTL8(FT0[i]);
        FT2[i] = ROTL8(FT1[i]);
        FT3[i] = ROTL8(FT2[i]);

        x = RSb[i];

        RT0[i] = ((unsigned long)MUL(0x0E, x)) ^
                 ((unsigned long)MUL(0x09, x) << 8) ^
                 ((unsigned long)MUL(0x0D, x) << 16) ^
                 ((unsigned long)MUL(0x0B, x) << 24);

        RT1[i] = ROTL8(RT0[i]);
        RT2[i] = ROTL8(RT1[i]);
        RT3[i] = ROTL8(RT2[i]);
    }

    if (tmp)
    {
        free(tmp);
        tmp = NULL;
    }
}

/*
 * AES key schedule (encryption)
 */
void OscAes_SetKey_Enc(StOscAesContext *ctx, unsigned char *key, int keysize)
{
    int i;
    unsigned long *RK;

    if (aes_init_done == 0)
    {
        OscAes_Gen_Tables();
        aes_init_done = 1;
    }

    switch (keysize)
    {
    case 128:
        ctx->nr = 10;
        break;
    case 192:
        ctx->nr = 12;
        break;
    case 256:
        ctx->nr = 14;
        break;
    default:
        return;
    }

    ctx->rk = RK = ctx->buf;

    for (i = 0; i < (keysize >> 5); i++)
        GET_ULONG_LE(RK[i], key, i << 2);

    switch (ctx->nr)
    {
    case 10:

        for (i = 0; i < 10; i++, RK += 4)
        {
            RK[4] = RK[0] ^ RCON[i] ^
                    ((unsigned long)FSb[(RK[3] >> 8) & 0xFF]) ^
                    ((unsigned long)FSb[(RK[3] >> 16) & 0xFF] << 8) ^
                    ((unsigned long)FSb[(RK[3] >> 24) & 0xFF] << 16) ^
                    ((unsigned long)FSb[(RK[3]) & 0xFF] << 24);

            RK[5] = RK[1] ^ RK[4];
            RK[6] = RK[2] ^ RK[5];
            RK[7] = RK[3] ^ RK[6];
        }
        break;

    case 12:

        for (i = 0; i < 8; i++, RK += 6)
        {
            RK[6] = RK[0] ^ RCON[i] ^
                    ((unsigned long)FSb[(RK[5] >> 8) & 0xFF]) ^
                    ((unsigned long)FSb[(RK[5] >> 16) & 0xFF] << 8) ^
                    ((unsigned long)FSb[(RK[5] >> 24) & 0xFF] << 16) ^
                    ((unsigned long)FSb[(RK[5]) & 0xFF] << 24);

            RK[7] = RK[1] ^ RK[6];
            RK[8] = RK[2] ^ RK[7];
            RK[9] = RK[3] ^ RK[8];
            RK[10] = RK[4] ^ RK[9];
            RK[11] = RK[5] ^ RK[10];
        }
        break;

    case 14:

        for (i = 0; i < 7; i++, RK += 8)
        {
            RK[8] = RK[0] ^ RCON[i] ^
                    ((unsigned long)FSb[(RK[7] >> 8) & 0xFF]) ^
                    ((unsigned long)FSb[(RK[7] >> 16) & 0xFF] << 8) ^
                    ((unsigned long)FSb[(RK[7] >> 24) & 0xFF] << 16) ^
                    ((unsigned long)FSb[(RK[7]) & 0xFF] << 24);

            RK[9] = RK[1] ^ RK[8];
            RK[10] = RK[2] ^ RK[9];
            RK[11] = RK[3] ^ RK[10];

            RK[12] = RK[4] ^
                     ((unsigned long)FSb[(RK[11]) & 0xFF]) ^
                     ((unsigned long)FSb[(RK[11] >> 8) & 0xFF] << 8) ^
                     ((unsigned long)FSb[(RK[11] >> 16) & 0xFF] << 16) ^
                     ((unsigned long)FSb[(RK[11] >> 24) & 0xFF] << 24);

            RK[13] = RK[5] ^ RK[12];
            RK[14] = RK[6] ^ RK[13];
            RK[15] = RK[7] ^ RK[14];
        }
        break;

    default:

        break;
    }
}

/*
 * AES key schedule (decryption)
 */
void OscAes_SetKey_Dec(StOscAesContext *ctx, unsigned char *key, int keysize)
{
    int i, j;
    StOscAesContext cty;
    unsigned long *RK;
    unsigned long *SK;

    switch (keysize)
    {
    case 128:
        ctx->nr = 10;
        break;
    case 192:
        ctx->nr = 12;
        break;
    case 256:
        ctx->nr = 14;
        break;
    default:
        return;
    }

    ctx->rk = RK = ctx->buf;

    OscAes_SetKey_Enc(&cty, key, keysize);
    SK = cty.rk + cty.nr * 4;

    *RK++ = *SK++;
    *RK++ = *SK++;
    *RK++ = *SK++;
    *RK++ = *SK++;

    for (i = ctx->nr - 1, SK -= 8; i > 0; i--, SK -= 8)
    {
        for (j = 0; j < 4; j++, SK++)
        {
            *RK++ = RT0[FSb[(*SK) & 0xFF]] ^
                    RT1[FSb[(*SK >> 8) & 0xFF]] ^
                    RT2[FSb[(*SK >> 16) & 0xFF]] ^
                    RT3[FSb[(*SK >> 24) & 0xFF]];
        }
    }

    *RK++ = *SK++;
    *RK++ = *SK++;
    *RK++ = *SK++;
    *RK++ = *SK++;

    memset(&cty, 0, sizeof(StOscAesContext));
}

#define AES_FROUND(X0, X1, X2, X3, Y0, Y1, Y2, Y3) \
    {                                              \
        X0 = *RK++ ^ FT0[(Y0) & 0xFF] ^            \
             FT1[(Y1 >> 8) & 0xFF] ^               \
             FT2[(Y2 >> 16) & 0xFF] ^              \
             FT3[(Y3 >> 24) & 0xFF];               \
                                                   \
        X1 = *RK++ ^ FT0[(Y1) & 0xFF] ^            \
             FT1[(Y2 >> 8) & 0xFF] ^               \
             FT2[(Y3 >> 16) & 0xFF] ^              \
             FT3[(Y0 >> 24) & 0xFF];               \
                                                   \
        X2 = *RK++ ^ FT0[(Y2) & 0xFF] ^            \
             FT1[(Y3 >> 8) & 0xFF] ^               \
             FT2[(Y0 >> 16) & 0xFF] ^              \
             FT3[(Y1 >> 24) & 0xFF];               \
                                                   \
        X3 = *RK++ ^ FT0[(Y3) & 0xFF] ^            \
             FT1[(Y0 >> 8) & 0xFF] ^               \
             FT2[(Y1 >> 16) & 0xFF] ^              \
             FT3[(Y2 >> 24) & 0xFF];               \
    }

#define AES_RROUND(X0, X1, X2, X3, Y0, Y1, Y2, Y3) \
    {                                              \
        X0 = *RK++ ^ RT0[(Y0) & 0xFF] ^            \
             RT1[(Y3 >> 8) & 0xFF] ^               \
             RT2[(Y2 >> 16) & 0xFF] ^              \
             RT3[(Y1 >> 24) & 0xFF];               \
                                                   \
        X1 = *RK++ ^ RT0[(Y1) & 0xFF] ^            \
             RT1[(Y0 >> 8) & 0xFF] ^               \
             RT2[(Y3 >> 16) & 0xFF] ^              \
             RT3[(Y2 >> 24) & 0xFF];               \
                                                   \
        X2 = *RK++ ^ RT0[(Y2) & 0xFF] ^            \
             RT1[(Y1 >> 8) & 0xFF] ^               \
             RT2[(Y0 >> 16) & 0xFF] ^              \
             RT3[(Y3 >> 24) & 0xFF];               \
                                                   \
        X3 = *RK++ ^ RT0[(Y3) & 0xFF] ^            \
             RT1[(Y2 >> 8) & 0xFF] ^               \
             RT2[(Y1 >> 16) & 0xFF] ^              \
             RT3[(Y0 >> 24) & 0xFF];               \
    }

/*
 * AES-ECB block encryption/decryption
 */
void OscAes_Crypt_Ecb(StOscAesContext *ctx,
                      int mode, unsigned char input[16], unsigned char output[16])
{
    int i;
    unsigned long *RK, X0, X1, X2, X3, Y0, Y1, Y2, Y3;

    RK = ctx->rk;

    GET_ULONG_LE(X0, input, 0);
    X0 ^= *RK++;
    GET_ULONG_LE(X1, input, 4);
    X1 ^= *RK++;
    GET_ULONG_LE(X2, input, 8);
    X2 ^= *RK++;
    GET_ULONG_LE(X3, input, 12);
    X3 ^= *RK++;

    if (mode == AES_DECRYPT)
    {
        for (i = (ctx->nr >> 1) - 1; i > 0; i--)
        {
            AES_RROUND(Y0, Y1, Y2, Y3, X0, X1, X2, X3);
            AES_RROUND(X0, X1, X2, X3, Y0, Y1, Y2, Y3);
        }

        AES_RROUND(Y0, Y1, Y2, Y3, X0, X1, X2, X3);

        X0 = *RK++ ^
             ((unsigned long)RSb[(Y0) & 0xFF]) ^
             ((unsigned long)RSb[(Y3 >> 8) & 0xFF] << 8) ^
             ((unsigned long)RSb[(Y2 >> 16) & 0xFF] << 16) ^
             ((unsigned long)RSb[(Y1 >> 24) & 0xFF] << 24);

        X1 = *RK++ ^
             ((unsigned long)RSb[(Y1) & 0xFF]) ^
             ((unsigned long)RSb[(Y0 >> 8) & 0xFF] << 8) ^
             ((unsigned long)RSb[(Y3 >> 16) & 0xFF] << 16) ^
             ((unsigned long)RSb[(Y2 >> 24) & 0xFF] << 24);

        X2 = *RK++ ^
             ((unsigned long)RSb[(Y2) & 0xFF]) ^
             ((unsigned long)RSb[(Y1 >> 8) & 0xFF] << 8) ^
             ((unsigned long)RSb[(Y0 >> 16) & 0xFF] << 16) ^
             ((unsigned long)RSb[(Y3 >> 24) & 0xFF] << 24);

        X3 = *RK++ ^
             ((unsigned long)RSb[(Y3) & 0xFF]) ^
             ((unsigned long)RSb[(Y2 >> 8) & 0xFF] << 8) ^
             ((unsigned long)RSb[(Y1 >> 16) & 0xFF] << 16) ^
             ((unsigned long)RSb[(Y0 >> 24) & 0xFF] << 24);
    }
    else
    {
        /* AES_ENCRYPT */
        for (i = (ctx->nr >> 1) - 1; i > 0; i--)
        {
            AES_FROUND(Y0, Y1, Y2, Y3, X0, X1, X2, X3);
            AES_FROUND(X0, X1, X2, X3, Y0, Y1, Y2, Y3);
        }

        AES_FROUND(Y0, Y1, Y2, Y3, X0, X1, X2, X3);

        X0 = *RK++ ^
             ((unsigned long)FSb[(Y0) & 0xFF]) ^
             ((unsigned long)FSb[(Y1 >> 8) & 0xFF] << 8) ^
             ((unsigned long)FSb[(Y2 >> 16) & 0xFF] << 16) ^
             ((unsigned long)FSb[(Y3 >> 24) & 0xFF] << 24);

        X1 = *RK++ ^
             ((unsigned long)FSb[(Y1) & 0xFF]) ^
             ((unsigned long)FSb[(Y2 >> 8) & 0xFF] << 8) ^
             ((unsigned long)FSb[(Y3 >> 16) & 0xFF] << 16) ^
             ((unsigned long)FSb[(Y0 >> 24) & 0xFF] << 24);

        X2 = *RK++ ^
             ((unsigned long)FSb[(Y2) & 0xFF]) ^
             ((unsigned long)FSb[(Y3 >> 8) & 0xFF] << 8) ^
             ((unsigned long)FSb[(Y0 >> 16) & 0xFF] << 16) ^
             ((unsigned long)FSb[(Y1 >> 24) & 0xFF] << 24);

        X3 = *RK++ ^
             ((unsigned long)FSb[(Y3) & 0xFF]) ^
             ((unsigned long)FSb[(Y0 >> 8) & 0xFF] << 8) ^
             ((unsigned long)FSb[(Y1 >> 16) & 0xFF] << 16) ^
             ((unsigned long)FSb[(Y2 >> 24) & 0xFF] << 24);
    }

    PUT_ULONG_LE(X0, output, 0);
    PUT_ULONG_LE(X1, output, 4);
    PUT_ULONG_LE(X2, output, 8);
    PUT_ULONG_LE(X3, output, 12);
}

/*
 * AES-CBC buffer encryption/decryption
 */
void OscAes_Crypt_Cbc(StOscAesContext *ctx, int mode, int length,
                      unsigned char iv[16], unsigned char *input, unsigned char *output, int *output_len)
{
    int i;
    unsigned char temp[16];
    *output_len = 0;
    if (mode == AES_DECRYPT)
    {
        while (length > 0)
        {
            memcpy(temp, input, 16);
            OscAes_Crypt_Ecb(ctx, mode, input, output);

            for (i = 0; i < 16; i++)
                output[i] = (unsigned char)(output[i] ^ iv[i]);

            memcpy(iv, temp, 16);

            input += 16;
            output += 16;
            length -= 16;
            *output_len += 16;
        }
    }
    else
    {
        while (length > 0)
        {
            for (i = 0; i < 16; i++)
                output[i] = (unsigned char)(input[i] ^ iv[i]);

            OscAes_Crypt_Ecb(ctx, mode, output, output);
            memcpy(iv, output, 16);

            input += 16;
            output += 16;
            length -= 16;
            *output_len += 16;
        }
    }
}

/*
 * AES-CFB128 buffer encryption/decryption
 */
void OscAes_Crypt_Cfb128(StOscAesContext *ctx,
                         int mode,
                         int length,
                         int *iv_off,
                         unsigned char iv[16],
                         unsigned char *input, unsigned char *output)
{
    int c, n = *iv_off;

    if (mode == AES_DECRYPT)
    {
        while (length--)
        {
            if (n == 0)
                OscAes_Crypt_Ecb(ctx, AES_ENCRYPT, iv, iv);

            c = *input++;
            *output++ = (unsigned char)(c ^ iv[n]);
            iv[n] = (unsigned char)c;

            n = (n + 1) & 0x0F;
        }
    }
    else
    {
        while (length--)
        {
            if (n == 0)
                OscAes_Crypt_Ecb(ctx, AES_ENCRYPT, iv, iv);

            iv[n] = *output++ = (unsigned char)(iv[n] ^ *input++);

            n = (n + 1) & 0x0F;
        }
    }

    *iv_off = n;
}

int OscAes_128Enc(char *aes_iv, char *aes_key, char *source, int source_len, char *dec, int *dec_len)
{
    StOscAesContext ctx;
    uint8_t iv[20] = {0};
    uint8_t private_key[36] = {0};
    int i = 0;

    strcpy((char *)iv, aes_iv);
    strcpy((char *)private_key, aes_key);
    int pad_byte_len = 16 - (source_len % 16);

    for (i = 0; i < pad_byte_len; i++)
    {
        source[source_len++] = pad_byte_len;
    }

    OscAes_SetKey_Enc(&ctx, (uint8_t *)private_key, 128);
    OscAes_Crypt_Cbc(&ctx, AES_ENCRYPT, source_len, iv, (unsigned char *)source, (unsigned char *)dec, dec_len);
    return 0;
}

int OscAes_128Dec(char *aes_iv, char *aes_key, char *source, int source_len, char *dec, int *dec_len)
{
    StOscAesContext ctx;
    unsigned char iv[BK_AES_IV_LEN + 1] = {0}; /*impor*/
    unsigned char private_key[BK_AES_KEY_LEN + 1] = {0};

    memcpy(iv, aes_iv, strlen(aes_iv));
    iv[sizeof(iv) - 1] = '\0';
    memcpy(private_key, aes_key, strlen(aes_key));
    private_key[sizeof(private_key) - 1] = '\0';

    printf("private_key %s\r\n", private_key);

    int enc_len = source_len;
    int supply_enc_len = source_len;
    unsigned char *supply_enc_data = (unsigned char *)malloc(supply_enc_len);

    memset(supply_enc_data, 0, supply_enc_len);
    memcpy(supply_enc_data, source, enc_len);
    memset((unsigned char *)dec, 0x0, *dec_len);
    OscAes_SetKey_Dec(&ctx, (unsigned char *)private_key, 128);

    OscAes_Crypt_Cbc(&ctx, AES_DECRYPT, supply_enc_len, iv, supply_enc_data, (unsigned char *)dec, dec_len);
    *dec_len -= dec[*dec_len - 1];
    dec[*dec_len] = '\0';
    free(supply_enc_data);

    return 0;
}

char *OscAes_128AndBase64_Encode(char *rxdata)
{
    char *aes_ecode = NULL;
    char *base64_ecode = NULL;

    aes_ecode = (char *)malloc(1024);
    if (aes_ecode == NULL)
    {
        printf("no memory for receive response buffer %s %d.\n", __FUNCTION__, __LINE__);
        goto __exit;
    }

    int len = 0;
    memset(aes_ecode, 0, 1024);
    OscAes_128Enc(DEF_AES_ECODE_IV, DEF_AES_ECODE_KEY, rxdata, strlen(rxdata), aes_ecode, &len);
    base64_ecode = (char *)OSC_base64_encode((unsigned char *)aes_ecode, len);

__exit:
    if (aes_ecode)
    {
        free(aes_ecode);
    }

    return base64_ecode;
}

/*
 * @param rxdata 待解码的数据
 * rxdata 用来存放编码后的json数据，会把传进来的数据覆盖；
 * @return NULL：编码失败；其他值：为解码数据的缓冲区；
 */
char *OscAes_128AndBase64_Decode(char *rxdata)
{
    char *base64_decode = NULL;
    char *aes_decode = NULL;

    base64_decode = (char *)malloc(1024);
    if (base64_decode == NULL)
    {
        printf("no memory for receive response buffer %s %d.\n", __FUNCTION__, __LINE__);
        goto __exit;
    }

    long base64_decode_len = 0;

    base64_decode = (char *)OSC_base64_decode_outlen((unsigned char *)rxdata, &base64_decode_len);
    printf("base64_decode_len :%ld \n", base64_decode_len);

    aes_decode = rxdata;
    memset(aes_decode, 0, strlen(aes_decode));
    if (aes_decode == NULL)
    {
        printf("no memory for receive response buffer %s %d.\n", __FUNCTION__, __LINE__);
        goto __exit;
    }

    int dec_len = strlen(aes_decode);
    OscAes_128Dec(DEF_AES_ECODE_IV, DEF_AES_ECODE_KEY, (char *)base64_decode, base64_decode_len, (char *)aes_decode, &dec_len);

__exit:
    if (base64_decode)
    {
        free(base64_decode);
    }

    return aes_decode;
}
